package com.example.exceldemo.write;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.example.exceldemo.entity.TUser;
import com.example.exceldemo.entity.UserExcel;
import com.example.exceldemo.mapper.TUserMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import java.util.Map;
import java.util.concurrent.*;

/**
 * @Author: yanhongwei
 * @Date: 2023/07/14 11:15:34
 * @Description: TODO
 */

@Slf4j
@Component
public class BatchExport {

    private TUserMapper userMapper;

    public BatchExport(TUserMapper userMapper) {
        this.userMapper = userMapper;
    }

    @Autowired
    @Qualifier("excelThreadPool")
    private ThreadPoolTaskExecutor threadPoolTaskExecutor;

    public void exportBatch() {

        // 每页多少条数据
        int pageSize = 10000;
        // 导出文件路径
        String fileName = "E:\\test.xlsx";
        // 查询条件
        QueryWrapper<TUser> queryWrapper = new QueryWrapper<>();
        Long start = System.currentTimeMillis();
        Long count = (long) userMapper.selectCount(queryWrapper);
        Long sheetNum = count % pageSize == 0 ? count / pageSize : count / pageSize + 1;
        // 多线程去读
        // 1.初始化map容量 防止扩容带来的效率损耗
        Map<Integer, Page<TUser>> pageMap = new ConcurrentHashMap<>(Math.toIntExact(sheetNum));
        CountDownLatch countDownLatch = new CountDownLatch(Math.toIntExact(sheetNum));
        // 注意 easyexcel 暂时不支持多线程并发写入！！！ 详情请看github上issues
        for (int i = 0; i < sheetNum; i++) {
            int finali = i;
            threadPoolTaskExecutor.submit(() -> {
                Page<TUser> page = new Page<>();
                page.setCurrent(finali + 1);
                page.setSize(pageSize);
                // 获取数据存放到map中
                Page<TUser> selectPage = userMapper.selectPage(page, queryWrapper);
                pageMap.put(finali, selectPage);
                // 消耗掉一个
                countDownLatch.countDown();
            });
        }
        try {
            countDownLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        // 写入
        try (ExcelWriter excelWriter = EasyExcel.write(fileName, TUser.class).build()) {
            pageMap.forEach((k, v) -> {
                log.info("正在写入{}条数据", pageSize);
                WriteSheet writeSheet = EasyExcel.writerSheet(k, "第" + (k + 1) + "批数据").build();
                excelWriter.write(v.getRecords(), writeSheet);
                // 写完当前数据立刻删除  不删除会产生内存泄漏即无法回收Map中巨大的空间 导致oom
                pageMap.remove(k);
            });
            excelWriter.finish();
        }
        Long end = System.currentTimeMillis();
        log.info("耗时" + (end - start) + "ms");
    }

}
